package org.apache.spark.streaming.rabbitmq

import java.io.IOException
import java.util.concurrent.{BlockingQueue, LinkedBlockingQueue}

import com.rabbitmq.client.{AMQP, Channel, ConsumerCancelledException, DefaultConsumer, Delivery, Envelope, ShutdownSignalException}
import com.rabbitmq.utility.Utility

/**
 * SparkSQL AMQP 读取消费者参数
 * <pre>
 *
 * Created by zhenqin.
 * User: zhenqin
 * Date: 2020/12/25
 * Time: 下午1:13
 * Vendor: yiidata.com
 *
 * </pre>
 *
 * @author zhenqin
 */
class QueueingConsumer(val channel: Channel, val capacity:Int) extends DefaultConsumer(channel) {


  // Marker object used to signal the queue is in shutdown mode.
  // It is only there to wake up consumers. The canonical representation
  // of shutting down is the presence of _shutdown.
  // Invariant: This is never on _queue unless _shutdown != null.
  private val POISON = new Delivery(null, null, null)
  private val _queue = new LinkedBlockingQueue[Delivery](capacity)
  // When this is non-null the queue is in shutdown mode and nextDelivery should
  // throw a shutdown signal exception.
  private var _shutdown: ShutdownSignalException = null
  private var _cancelled: ConsumerCancelledException = null


  /**
   * 默认提供 10000 条记录的容量，方式内存过量使用
   * @param channel
   */
  def this(channel:Channel) {
    this(channel, 10000)
  }


  @throws[InterruptedException]
  @throws[ShutdownSignalException]
  @throws[ConsumerCancelledException]
  def nextDelivery(): Delivery = handle(_queue.take)

  override def handleShutdownSignal(consumerTag: String, sig: ShutdownSignalException): Unit = {
    _shutdown = sig
    _queue.add(POISON)
  }

  @throws[IOException]
  override def handleCancel(consumerTag: String): Unit = {
    _cancelled = new ConsumerCancelledException
    _queue.add(POISON)
  }

  @throws[IOException]
  override def handleDelivery(consumerTag: String, envelope: Envelope, properties: AMQP.BasicProperties, body: Array[Byte]): Unit = {
    checkShutdown()
    this._queue.add(new Delivery(envelope, properties, body))
  }

  /**
   * Check if we are in shutdown mode and if so throw an exception.
   */
  private def checkShutdown(): Unit = {
    if (_shutdown != null) throw Utility.fixStackTrace(_shutdown)
  }

  /**
   * If delivery is not POISON nor null, return it.
   * <p/>
   * If delivery, _shutdown and _cancelled are all null, return null.
   * <p/>
   * If delivery is POISON re-insert POISON into the queue and
   * throw an exception if POISONed for no reason.
   * <p/>
   * Otherwise, if we are in shutdown mode or cancelled,
   * throw a corresponding exception.
   */
  private def handle(delivery: Delivery) = {
    if ((delivery eq POISON) || delivery == null && (_shutdown != null || _cancelled != null)) {
      if (delivery eq POISON) {
        _queue.add(POISON)
        if (_shutdown == null && _cancelled == null) throw new IllegalStateException("POISON in queue, but null _shutdown and null _cancelled. " + "This should never happen, please report as a BUG")
      }
      if (null != _shutdown) throw Utility.fixStackTrace(_shutdown)
      if (null != _cancelled) throw Utility.fixStackTrace(_cancelled)
    }
    delivery
  }
}
